package main

import (
	"encoding/json"
	"fmt"
	"github.com/samuel/go-zookeeper/zk"
	"log"
	"time"
)

type ServiceNode struct {
	Name string `json:"name"`
	Host string `json:"host"`
	Port int    `json:"post"`
}

type SdClient struct {
	//服务列表
	zkServers []string
	//服务根目录
	zkRoot string

	conn *zk.Conn
}

func NewClient(zkServers []string, zkRoot string) (*SdClient, error) {
	//创建客户端
	client := new(SdClient)
	client.zkServers = zkServers
	client.zkRoot = zkRoot

	//连接服务器
	conn, _, err := zk.Connect(zkServers, time.Second*50)
	if err != nil {
		return nil, err
	}

	client.conn = conn

	//创建服务根节点
	if err = client.createRootNode(); err != nil {
		//关闭连接
		client.Close()
		return nil, err
	}

	return client, nil
}

func (s *SdClient) createRootNode() error {
	//判断根节点是否存在
	exists, _, err := s.conn.Exists(s.zkRoot)
	if err != nil {
		return err
	}

	//根节点不存在，则创建根节点
	if !exists {
		//创建持久化节点为根节点
		_, err := s.conn.Create(s.zkRoot, []byte(""), 0, zk.WorldACL(zk.PermAll))
		//集群多线程可能会导致create返回已存在的错误
		if err != nil && err != zk.ErrNodeExists {
			return err
		}
	}
	return nil
}

//关闭连接  释放临时节点
func (s *SdClient) Close() {
	s.conn.Close()
}

var hosts = []string{"node01:2181", "node02:2181", "node03:2181"}

func main() {

	client, err := NewClient(hosts, "/api")
	if err != nil {
		log.Fatal(err)
	}

	//模拟zk注册服务
	node01 := &ServiceNode{"user", "node01", 4001}
	node02 := &ServiceNode{"user", "node02", 4002}
	node03 := &ServiceNode{"user", "node03", 4003}

	//将3个模拟的节点注册到zk
	if err := client.register(node01);err!=nil{
		log.Fatal(err)
	}

	if err := client.register(node02);err!=nil{
		log.Fatal(err)
	}
	if err := client.register(node03);err!=nil{
		log.Fatal(err)
	}

	//服务消费 将三个注册的节点打印出来
	nodes, err := client.consumer("/api", "/user")
	for _,n := range nodes  {
		fmt.Printf("name:%s host: %s port: %d \n",n.Name,n.Host,n.Port)
	}

}

//模拟消费者
func (s *SdClient) consumer(zkRoot string, name string) ([] *ServiceNode, error){
	//定义路径 /api/user
	path := zkRoot + name
	//获取子节点名称
	children, _, err := s.conn.Children(path)
	if err != nil {
		if err == zk.ErrNoNode{
			return [] *ServiceNode{}, nil
		}
		return nil,err
	}

	nodes:= []*ServiceNode{}
	for _,child:= range children {
		fullPath := path+"/"+child
		data, _, err := s.conn.Get(fullPath)
		if err != nil {
			if err == zk.ErrNoNode {
				continue
			}
			return nil,err
		}
		node := new(ServiceNode)
		err = json.Unmarshal(data, node)
		if err != nil {
			return nil,err
		}
		nodes = append(nodes, node)
	}
	return nodes,nil
}

func (s *SdClient) register(node *ServiceNode) error {
	//创建api/user目录
	if err := s.ensureName(node.Name); err != nil {
		return err
	}

	path:= s.zkRoot+"/"+node.Name+"/n"
	//模拟给节点设置数据，将node转为数组
	data, err := json.Marshal(node)
	if err != nil {
		return err
	}
	_, err = s.conn.CreateProtectedEphemeralSequential(path, data, zk.WorldACL(zk.PermAll))
	if err != nil {
		return err
	}
	return nil

}

func (s *SdClient) ensureName(name string) error {
	//定义检查的目录/api/user
	path := s.zkRoot + "/" + name
	//验证目录是否存在
	bol, _, err := s.conn.Exists(path)
	if err != nil && err == zk.ErrNodeExists {
		return err
	}

	if !bol {
		_, err := s.conn.Create(path, []byte(""), 0, zk.WorldACL(zk.PermAll))
		if err != nil && err == zk.ErrNodeExists {
			return err
		}
	}
	return nil
}
